package GMDPRE

import (
	"fmt"
	"math/rand"
	"time"
)

const DEBUG = true
const Threshold = 2
const NumOfFrag = 3

var sysRand = rand.New(rand.NewSource(time.Now().UnixNano()))

type System struct {
	N      int
	t      int
	PNs    [NumOfFrag]ProxyNode
	idOfPN [NumOfFrag][]uint8
	random *rand.Rand
}

var sys System

func (s *System) Init() {
	sys.N = NumOfFrag
	sys.t = Threshold

	for i := 0; i < NumOfFrag; i++ {
		sys.PNs[i].Init()
		sys.idOfPN[i] = BnToBytes(sys.PNs[i].GetId(), Sm2CurveDigit)
	}
}

func (s *System) ReEnCapsulate(KF *[NumOfFrag][]uint8, capBytes []uint8, PubKeyBytes []uint8) [][]uint8 {
	var res [][]uint8
	for i := 0; i < NumOfFrag; i++ {
		res = append(res, s.PNs[i].ReEnCapsulate(KF[i], capBytes, PubKeyBytes))
	}
	return res
}

func UserKeyGen(u *User) {
	u.Init()
}

func TestAPP() {
	var ret int

	sys.Init()
	//BenchStart()
	UserKeyGen(&UserA)
	UserKeyGen(&UserB)
	//BenchEnd()

	fmt.Println("+---------------------------------用户注册阶段----------------------------------------+")
	fmt.Println("数据持有者A的私钥: ")
	//PrintBytesHex(BnToBytes(UserA.GetId(), Sm2CurveDigit), Sm2CurveSize)
	PrintBytesHex(BnToBytes(UserA.getPriKey(), Sm2CurveDigit), Sm2CurveSize)
	fmt.Println("数据持有者A的公钥: ")
	PrintBytesHex(UserA.GetPubKeyBytes(), Sm2AffPointSize)
	fmt.Println("数据需求者B的私钥: ")
	//PrintBytesHex(BnToBytes(UserB.GetId(), Sm2CurveDigit), Sm2CurveSize)
	PrintBytesHex(BnToBytes(UserB.getPriKey(), Sm2CurveDigit), Sm2CurveSize)
	fmt.Println("数据持有者B的公钥: ")
	PrintBytesHex(UserB.GetPubKeyBytes(), Sm2AffPointSize)
	//fmt.Printf("A和B注册消耗时间: %.2f ms\n", CostTime())
	fmt.Println("+-------------------------------------------------------------------------------------+")

	fmt.Println("+-----------------------------加密与密钥封装阶段--------------------------------------+")
	//BenchStart()
	var capsule = UserA.EnCapsulate()
	//BenchEnd()
	fmt.Println("A生成的原始加密密钥:")
	PrintBytesHex(UserA.KO, Sm2AffPointSize)
	fmt.Println("A生成的密钥封装:")
	PrintBytesHex(capsule, CapsuleSize)
	//fmt.Printf("密钥封装消耗时间: %.2f ms\n", CostTime())

	var sm4Key []uint8
	var C []uint8
	var M = []uint8{
		0x4d, 0x61, 0x6b, 0x69, 0x6e, 0x67, 0x20, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x73, 0x20, 0x69, 0x6e,
		0x20, 0x4c, 0x61, 0x54, 0x65, 0x58, 0x20, 0x63, 0x61, 0x6e, 0x20, 0x62, 0x65, 0x20, 0x74, 0x65,
		0x64, 0x69, 0x6f, 0x75, 0x73, 0x2c, 0x20, 0x65, 0x73, 0x70, 0x65, 0x63, 0x69, 0x61, 0x6c, 0x6c,
		0x79, 0x20, 0x69, 0x66, 0x20, 0x73, 0x6f, 0x6d, 0x65, 0x20, 0x63, 0x6f, 0x6c, 0x75, 0x6d, 0x6e}
	var sig []uint8

	//BenchStart()
	sm4Key = KDF(Sm4BlockSize, UserA.KO, Sm2AffPointSize)
	C = Sm4Encrypt(M, len(M), sm4Key)
	for {
		rnd := GetRandomBytes(32)
		ret, sig = Sm2SignHash(BnToBytes(UserA.getPriKey(), Sm2CurveDigit), rnd, Sm3Hash(M, uint64(len(M))))
		if ret == 0 {
			break
		}
	}
	//BenchEnd()

	fmt.Print("A扩展出的对称密钥: ")
	PrintBytesHex(sm4Key, Sm4BlockSize)
	fmt.Println("A要加密的明文:")
	PrintBytesHex(M, len(M))
	fmt.Println("加密后的密文:")
	PrintBytesHex(C, len(C))
	fmt.Println("明文的签名: ")
	PrintBytesHex(sig, Sm2CurveSize*2)
	//fmt.Printf("加密消耗时间: %.2f ms\n", CostTime())
	fmt.Println("+-------------------------------------------------------------------------------------+")

	fmt.Println("+-----------------------------重加密密钥生成阶段--------------------------------------+")
	var KF [NumOfFrag][]uint8

	BenchStart()
	UserA.ReKeyGen(&KF, UserB.GetPubKeyBytes(), &sys.idOfPN)
	BenchEnd()

	for i := 0; i < NumOfFrag; i++ {
		fmt.Printf("A为代理重加密节点%d生成的重加密密钥片段:\n", i)
		PrintBytesHex(KF[i], FragSize)
	}
	//fmt.Printf("重加密密钥生成消耗时间: %.2f ms\n", CostTime())
	fmt.Println("+-------------------------------------------------------------------------------------+")

	fmt.Println("+----------------------------------重加密阶段-----------------------------------------+")
	//KF[0][0] ^= 1
	//BenchStart()
	var cF = sys.ReEnCapsulate(&KF, capsule, UserA.GetPubKeyBytes())
	//BenchEnd()

	for i := 0; i < NumOfFrag; i++ {
		fmt.Printf("代理重加密节点%d为B生成的重加密密文片段:\n", i)
		PrintBytesHex(cF[i], CFragSize)
	}
	//fmt.Printf("重加密消耗时间: %.2f ms\n", CostTime()/float64(NumOfFrag))
	fmt.Println("+-------------------------------------------------------------------------------------+")

	fmt.Println("+----------------------------------解密1阶段------------------------------------------+")
	//BenchStart()
	var KO = UserB.DeCapsulateFrags(UserA.GetPubKeyBytes(), cF)
	//BenchEnd()

	fmt.Println("B还原出来的原始加密密钥:")
	PrintBytesHex(KO, Sm2AffPointSize)
	//fmt.Printf("原始密钥恢复消耗时间: %.2f ms\n", CostTime())

	//BenchStart()
	sm4Key = KDF(Sm4BlockSize, KO, Sm2AffPointSize)
	var CN = Sm4Decrypt(C, len(C), sm4Key)
	var v = Sm2VrfHashSign(Sm3Hash(CN, uint64(len(CN))), UserA.GetPubKeyBytes(), sig)
	//BenchEnd()

	fmt.Print("B扩展出的对称密钥: ")
	PrintBytesHex(sm4Key, Sm4BlockSize)
	fmt.Printf("B解密是否成功: %v\n", v == 0)
	fmt.Println("B解密出的明文:")
	PrintBytesHex(CN, len(CN))
	//fmt.Printf("解密消耗时间: %.2f ms\n", CostTime())
	fmt.Println("+-------------------------------------------------------------------------------------+")

	fmt.Println("+----------------------------------解密2阶段------------------------------------------+")
	//BenchStart()
	var a = UserA.DeCapsulate(capsule)
	//BenchEnd()

	fmt.Println("A还原出来的原始加密密钥:")
	PrintBytesHex(a, Sm2AffPointSize)
	//fmt.Printf("原始密钥恢复消耗时间: %.2f ms\n", CostTime())

	//BenchStart()
	sm4Key = KDF(Sm4BlockSize, KO, Sm2AffPointSize)
	var MN = Sm4Decrypt(C, len(C), sm4Key)
	v = Sm2VrfHashSign(Sm3Hash(MN, uint64(len(MN))), UserA.GetPubKeyBytes(), sig)
	//BenchEnd()

	fmt.Print("A扩展出的对称密钥: ")
	PrintBytesHex(sm4Key, Sm4BlockSize)
	fmt.Printf("A解密是否成功: %v\n", v == 0)
	fmt.Println("A解密出的明文:")
	PrintBytesHex(CN, len(CN))
	//fmt.Printf("解密消耗时间: %.2f ms\n", CostTime())
	fmt.Println("+-------------------------------------------------------------------------------------+")
}

func testAPPV2() (float64, float64, float64, float64, float64, float64) {
	var ret int

	sys.Init()

	//fmt.Println("+---------------------------------用户注册阶段----------------------------------------+")
	UserKeyGen(&UserA)
	BenchStart()
	UserKeyGen(&UserB)
	BenchEnd()
	t1 := CostTime()
	//fmt.Println("+-------------------------------------------------------------------------------------+")

	//fmt.Println("+-----------------------------加密与密钥封装阶段--------------------------------------+")
	var sm4Key []uint8
	var C []uint8
	//var M = []uint8{
	//	0x4d, 0x61, 0x6b, 0x69, 0x6e, 0x67, 0x20, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x73, 0x20, 0x69, 0x6e,
	//	0x20, 0x4c, 0x61, 0x54, 0x65, 0x58, 0x20, 0x63, 0x61, 0x6e, 0x20, 0x62, 0x65, 0x20, 0x74, 0x65,
	//	0x64, 0x69, 0x6f, 0x75, 0x73, 0x2c, 0x20, 0x65, 0x73, 0x70, 0x65, 0x63, 0x69, 0x61, 0x6c, 0x6c,
	//	0x79, 0x20, 0x69, 0x66, 0x20, 0x73, 0x6f, 0x6d, 0x65, 0x20, 0x63, 0x6f, 0x6c, 0x75, 0x6d, 0x6e}
	var M = GetRandomBytes(1024 * 1024)
	var sig []uint8

	BenchStart()
	var capsule = UserA.EnCapsulate()
	sm4Key = KDF(Sm4BlockSize, UserA.KO, Sm2AffPointSize)
	C = Sm4Encrypt(M, len(M), sm4Key)
	for {
		rnd := GetRandomBytes(32)
		ret, sig = Sm2SignHash(BnToBytes(UserA.getPriKey(), Sm2CurveDigit), rnd, Sm3Hash(M, uint64(len(M))))
		if ret == 0 {
			break
		}
	}
	BenchEnd()
	t2 := CostTime()
	//fmt.Println("+-------------------------------------------------------------------------------------+")

	//fmt.Println("+-----------------------------重加密密钥生成阶段--------------------------------------+")
	var KF [NumOfFrag][]uint8

	BenchStart()
	UserA.ReKeyGen(&KF, UserB.GetPubKeyBytes(), &sys.idOfPN)
	BenchEnd()
	t3 := CostTime()
	//fmt.Println("+-------------------------------------------------------------------------------------+")

	//fmt.Println("+----------------------------------重加密阶段-----------------------------------------+")
	//KF[0][0] ^= 1
	BenchStart()
	var cF = sys.ReEnCapsulate(&KF, capsule, UserA.GetPubKeyBytes())
	BenchEnd()
	t4 := CostTime() / float64(NumOfFrag)
	//fmt.Println("+-------------------------------------------------------------------------------------+")

	//fmt.Println("+----------------------------------解密1阶段------------------------------------------+")
	BenchStart()
	var KO = UserB.DeCapsulateFrags(UserA.GetPubKeyBytes(), cF)
	sm4Key = KDF(Sm4BlockSize, KO, Sm2AffPointSize)
	var CN = Sm4Decrypt(C, len(C), sm4Key)
	var v = Sm2VrfHashSign(Sm3Hash(CN, uint64(len(CN))), UserA.GetPubKeyBytes(), sig)
	BenchEnd()

	if v != 0 {
		fmt.Println("EEEEEEEEEEEEEEEEEEEEEEEEEEEEEE!!!!!!!!!!!!!!!!!!!")
	}

	t5 := CostTime()
	//fmt.Println("+-------------------------------------------------------------------------------------+")

	//fmt.Println("+----------------------------------解密2阶段------------------------------------------+")
	BenchStart()
	UserA.DeCapsulate(capsule)
	sm4Key = KDF(Sm4BlockSize, KO, Sm2AffPointSize)
	var MN = Sm4Decrypt(C, len(C), sm4Key)
	v = Sm2VrfHashSign(Sm3Hash(MN, uint64(len(MN))), UserA.GetPubKeyBytes(), sig)
	BenchEnd()

	t6 := CostTime()
	//fmt.Println("+-------------------------------------------------------------------------------------+")
	return t1, t2, t3, t4, t5, t6
}

func BenchAPP() {
	var times = 1000
	var t1, t2, t3, t4, t5, t6 float64
	var tt1, tt2, tt3, tt4, tt5, tt6 float64
	for i := 0; i < times; i++ {
		tt1, tt2, tt3, tt4, tt5, tt6 = testAPPV2()
		t1 += tt1
		t2 += tt2
		t3 += tt3
		t4 += tt4
		t5 += tt5
		t6 += tt6
	}
	fmt.Printf("%.2f\n", t1/float64(times))
	fmt.Printf("%.2f\n", t2/float64(times))
	fmt.Printf("%.2f\n", t3/float64(times))
	fmt.Printf("%.2f\n", t4/float64(times))
	fmt.Printf("%.2f\n", t5/float64(times))
	fmt.Printf("%.2f\n", t6/float64(times))
}
